"""
ChromaDB向量数据库适配器
提供与现有VectorDB接口兼容的ChromaDB实现
支持回退到现有向量数据库
"""

import os
import json
from typing import List, Dict, Any, Optional, Tuple
import numpy as np

# 尝试导入ChromaDB，失败时使用回退
try:
    import chromadb
    HAS_CHROMADB = True
except ImportError:
    HAS_CHROMADB = False
    print("⚠️  ChromaDB未安装，相关功能将使用现有向量数据库")

# 导入现有向量数据库用于回退
from .vector_db import VectorDB, KnowledgeVectorDB


class ChromaDB:
    """ChromaDB向量数据库实现"""
    
    def __init__(self, collection_name: str = "knowledge_base", persist_directory: str = "./chroma_db"):
        """初始化ChromaDB客户端"""
        self.collection_name = collection_name
        self.persist_directory = persist_directory
        
        if HAS_CHROMADB:
            self.client = chromadb.Client()
            # 创建或获取集合
            self.collection = self.client.get_or_create_collection(
                name=collection_name,
                metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
            )
            self.mode = "chromadb"
        else:
            # 使用现有向量数据库回退
            self.vector_db = VectorDB()
            self.mode = "fallback"
            print(f"⚠️  ChromaDB未安装，使用现有向量数据库回退模式")
        
        self.documents = []
        self.metadata = []
    
    def add_documents(self, documents: List[str], metadatas: Optional[List[Dict]] = None):
        """添加文档到集合"""
        if not documents:
            return
            
        if self.mode == "chromadb":
            # 生成文档ID
            doc_ids = [f"doc_{i}" for i in range(len(self.documents), len(self.documents) + len(documents))]
            
            # 添加文档到内存缓存
            self.documents.extend(documents)
            if metadatas:
                self.metadata.extend(metadatas)
            else:
                self.metadata.extend([{}] * len(documents))
            
            # 添加到ChromaDB
            self.collection.add(
                documents=documents,
                metadatas=metadatas or [{}] * len(documents),
                ids=doc_ids
            )
        else:
            # 使用回退模式
            self.vector_db.add_documents(documents, metadatas)
    
    def search(self, query: str, n_results: int = 5) -> List[Tuple[str, float, Dict]]:
        """搜索相似文档"""
        if not query.strip():
            return []
        
        if self.mode == "chromadb":
            # 使用ChromaDB搜索
            results = self.collection.query(
                query_texts=[query],
                n_results=n_results
            )
            
            # 格式化结果
            formatted_results = []
            if results['documents']:
                for i, doc in enumerate(results['documents'][0]):
                    score = results['distances'][0][i] if results['distances'] else 0.0
                    metadata = results['metadatas'][0][i] if results['metadatas'] else {}
                    formatted_results.append((doc, score, metadata))
            
            return formatted_results
        else:
            # 使用回退模式
            return self.vector_db.search(query, n_results)
    
    def get_document_count(self) -> int:
        """获取文档数量"""
        if self.mode == "chromadb":
            return self.collection.count()
        else:
            return self.vector_db.get_document_count()
    
    def clear(self):
        """清空集合"""
        if self.mode == "chromadb":
            self.client.delete_collection(self.collection_name)
            self.collection = self.client.get_or_create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": "cosine"}
            )
        else:
            self.vector_db.clear()
        
        self.documents = []
        self.metadata = []
    
    def save(self, file_path: str):
        """保存向量数据库状态"""
        # ChromaDB会自动持久化，此方法为兼容接口
        pass
    
    def load(self, file_path: str):
        """加载向量数据库状态"""
        # ChromaDB会自动加载，此方法为兼容接口
        pass


class KnowledgeChromaDB:
    """知识库专用的ChromaDB管理器"""
    
    def __init__(self, role_name: str):
        """初始化指定角色的向量数据库"""
        self.role_name = role_name
        
        if HAS_CHROMADB:
            try:
                import chromadb
                from chromadb.config import Settings
                
                self.mode = "chromadb"
                self.client = chromadb.Client(Settings(
                    chroma_db_impl="duckdb+parquet",
                    persist_directory="./chroma_db"
                ))
                
                # 创建或获取集合
                self.collection = self.client.get_or_create_collection(
                    name=f"knowledge_{role_name}",
                    metadata={"hnsw:space": "cosine"}
                )
                
                print(f"✅ {role_name}: 知识库ChromaDB初始化完成")
                
            except Exception as e:
                print(f"⚠️  ChromaDB初始化失败: {e}")
                print("⚠️  将使用回退方案")
                self._init_fallback(role_name)
        else:
            self._init_fallback(role_name)
    
    def add_qa_pair(self, question: str, answer: str, metadata: Optional[Dict] = None):
        """添加问答对到向量数据库"""
        if self.mode == "chromadb":
            # 组合问题和答案作为文档
            document = f"问题: {question}\n答案: {answer}"
            
            # 准备元数据
            qa_metadata = {
                "question": question,
                "answer": answer,
                "role": self.role_name,
                "type": "qa_pair"
            }
            
            if metadata:
                qa_metadata.update(metadata)
            
            self.chroma_db.add_documents([document], [qa_metadata])
        else:
            # 使用回退模式 - 需要先保存到知识库文件，然后更新向量数据库
            kb_path = os.path.join(self.knowledge_base_dir, f"{self.role_name}.json")
            
            # 加载或创建知识库文件
            kb_data = {}
            if os.path.exists(kb_path):
                try:
                    with open(kb_path, 'r', encoding='utf-8') as f:
                        kb_data = json.load(f)
                except:
                    kb_data = {'qa_pairs': []}
            else:
                kb_data = {'qa_pairs': []}
            
            # 添加问答对
            if 'qa_pairs' not in kb_data:
                kb_data['qa_pairs'] = []
            
            qa_pair = {'question': question, 'answer': answer}
            if metadata:
                qa_pair['metadata'] = metadata
            
            kb_data['qa_pairs'].append(qa_pair)
            
            # 保存知识库文件
            with open(kb_path, 'w', encoding='utf-8') as f:
                json.dump(kb_data, f, ensure_ascii=False, indent=2)
            
            # 更新向量数据库
            self.vector_db.update_knowledge_base(self.role_name)
    
    def search_question(self, query: str, n_results: int = 3) -> List[Tuple[str, str, float]]:
        """搜索相关问题"""
        if self.mode == "chromadb":
            results = self.chroma_db.search(query, n_results)
            
            formatted_results = []
            for doc, score, metadata in results:
                # 从元数据中提取问题和答案
                question = metadata.get('question', '')
                answer = metadata.get('answer', '')
                
                # 如果元数据中没有，尝试从文档中解析
                if not question or not answer:
                    if "问题:" in doc and "答案:" in doc:
                        parts = doc.split("答案:")
                        if len(parts) >= 2:
                            question = parts[0].replace("问题:", "").strip()
                            answer = parts[1].strip()
                
                if question and answer:
                    formatted_results.append((question, answer, score))
            
            return formatted_results
        else:
            # 使用回退模式
            results = self.vector_db.search(self.role_name, query, top_k=n_results, threshold=0.1)
            
            formatted_results = []
            for doc, score, metadata in results:
                # 从元数据中提取问题和答案
                question = metadata.get('question', '')
                answer = metadata.get('answer', '')
                
                # 如果元数据中没有，使用文档内容
                if not question:
                    question = doc
                if not answer:
                    answer = metadata.get('answer', '未知答案')
                
                formatted_results.append((question, answer, score))
            
            return formatted_results
    
    def batch_add_qa_pairs(self, qa_pairs: List[Dict]):
        """批量添加问答对"""
        if self.mode == "chromadb":
            documents = []
            metadatas = []
            
            for qa in qa_pairs:
                question = qa.get('question', '')
                answer = qa.get('answer', '')
                
                if question and answer:
                    document = f"问题: {question}\n答案: {answer}"
                    metadata = {
                        "question": question,
                        "answer": answer,
                        "role": self.role_name,
                        "type": "qa_pair"
                    }
                    
                    # 合并额外元数据
                    if 'metadata' in qa and isinstance(qa['metadata'], dict):
                        metadata.update(qa['metadata'])
                    
                    documents.append(document)
                    metadatas.append(metadata)
            
            if documents:
                self.chroma_db.add_documents(documents, metadatas)
        else:
            # 使用回退模式
            kb_path = os.path.join(self.knowledge_base_dir, f"{self.role_name}.json")
            
            # 加载或创建知识库文件
            kb_data = {}
            if os.path.exists(kb_path):
                try:
                    with open(kb_path, 'r', encoding='utf-8') as f:
                        kb_data = json.load(f)
                except:
                    kb_data = {'qa_pairs': []}
            else:
                kb_data = {'qa_pairs': []}
            
            # 添加问答对
            if 'qa_pairs' not in kb_data:
                kb_data['qa_pairs'] = []
            
            for qa in qa_pairs:
                question = qa.get('question', '')
                answer = qa.get('answer', '')
                
                if question and answer:
                    qa_pair = {'question': question, 'answer': answer}
                    metadata = qa.get('metadata', {})
                    if metadata:
                        qa_pair['metadata'] = metadata
                    
                    kb_data['qa_pairs'].append(qa_pair)
            
            # 保存知识库文件
            with open(kb_path, 'w', encoding='utf-8') as f:
                json.dump(kb_data, f, ensure_ascii=False, indent=2)
            
            # 更新向量数据库
            self.vector_db.update_knowledge_base(self.role_name)
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        if self.mode == "chromadb":
            return {
                "role": self.role_name,
                "document_count": self.chroma_db.get_document_count(),
                "mode": "chromadb"
            }
        else:
            # 获取文档数量
            doc_count = 0
            if self.role_name in self.vector_db.vector_dbs:
                doc_count = self.vector_db.vector_dbs[self.role_name].size()
            
            return {
                "role": self.role_name,
                "document_count": doc_count,
                "mode": "fallback"
            }


def test_chroma_db():
    """测试ChromaDB功能"""
    print("=== 测试ChromaDB功能 ===")
    
    # 创建测试知识库
    kb = KnowledgeChromaDB("test_role")
    
    # 添加示例问答对
    test_qa_pairs = [
        {"question": "你好", "answer": "你好！我是AI助手"},
        {"question": "你是谁", "answer": "我是一个智能对话助手"},
        {"question": "帮助", "answer": "我可以回答问题和提供帮助"}
    ]
    
    kb.batch_add_qa_pairs(test_qa_pairs)
    print(f"添加了 {kb.get_stats()['document_count']} 个问答对")
    
    # 测试搜索
    test_queries = ["你好", "帮助我", "介绍你自己"]
    for query in test_queries:
        results = kb.search_question(query)
        print(f"搜索 '{query}':")
        for question, answer, score in results:
            print(f"  - {question} -> {answer} (相似度: {score:.3f})")
    
    print("=== 测试完成 ===")


if __name__ == "__main__":
    test_chroma_db()